package com.xiaoniu.kafkastreaming

import kafka.common.TopicAndPartition
import org.apache.kafka.common.serialization.StringDeserializer
import org.apache.spark.streaming.kafka.KafkaCluster
import org.apache.spark.streaming.kafka010.ConsumerStrategies.Subscribe
import org.apache.spark.streaming.kafka010.{CanCommitOffsets, HasOffsetRanges, KafkaUtils, OffsetRange}
import org.apache.spark.streaming.kafka010.LocationStrategies.PreferConsistent
import org.apache.spark.{SparkConf, SparkContext, TaskContext}
import org.apache.spark.streaming.{Seconds, StreamingContext}

/**
  * Created by xn066987 on 2017/7/24.
  */
object KafkaDirectStream {

  val brokerAddress = "10.17.2.110:9092"
  val groupId = "KafkaDirectStream11"
  val topic = "xnsparkstreaming11"
  val kafkaParams = scala.Predef.Map[String, Object](
    "bootstrap.servers" -> brokerAddress,
    "key.deserializer" -> classOf[StringDeserializer],
    "value.deserializer" -> classOf[StringDeserializer],
    "group.id" -> groupId,
    "auto.offset.reset" -> "latest",
    "enable.auto.commit" -> (false: java.lang.Boolean)
  )

  val kafkaStringParams = scala.Predef.Map[String, String](
    "bootstrap.servers" -> brokerAddress,
    "group.id" -> groupId
  )

  val topics = Array(topic)
  val master = "local[2]"

  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().setAppName("KafkaDirectStream11")
    sparkConf.setMaster(master)
    val sc = new SparkContext(sparkConf)
    sc.setLogLevel("WARN")
    val ssc = new StreamingContext(sc, Seconds(5))
    val messages = KafkaUtils.createDirectStream[String, String](
      ssc,
      PreferConsistent,
      Subscribe[String, String](topics, kafkaParams)
    )
    messages.map(record => (record.key, record.value)).flatMap(tuple => tuple._2.split(" ")).map { word => (word, 1) }.reduceByKey(_ + _).print()

    val kc = new KafkaCluster(kafkaStringParams)
    messages.foreachRDD(rdd => {
      val offsetRanges = rdd.asInstanceOf[HasOffsetRanges].offsetRanges
      rdd.foreachPartition { iter =>
        val offsets: OffsetRange = offsetRanges(TaskContext.get.partitionId)
        val topicAndPartition = TopicAndPartition(topic, offsets.partition)
        val result = kc.setConsumerOffsets(groupId, scala.collection.immutable.Map((topicAndPartition, offsets.untilOffset)))
        if (result.isLeft) {
          println(s"Error updating the offset to Kafka cluster: ${result.left.get}")
        }
      }
      // some time later, after outputs have completed
      messages.asInstanceOf[CanCommitOffsets].commitAsync(offsetRanges)
    })
    ssc.start()
    ssc.awaitTermination()
    ssc.stop()
  }

}
